use std::env;
use std::fs::OpenOptions;
use std::io::{self, Write};
use std::process::{Command, Stdio};

mod deepseek;
use deepseek::ask_deepseek;

fn main() {
    loop {
        print_prompt();
        io::stdout().flush().unwrap();
        
        let mut input = String::new();
        if io::stdin().read_line(&mut input).is_err() {
            eprintln!("Error reading input");
            continue;
        }
        
        let input = input.trim();
        if input.is_empty() {
            continue;
        }
        
        // 检查退出命令
        if input == "exit" {
            break;
        }
        
        // 处理问号开头的特殊命令
        if input.starts_with("??") {
            // 两个问号，使用 R1 大模型
            handle_deepseek_command(&input[2..].trim(), "r1");
            continue;
        } else if input.starts_with('?') {
            // 单个问号，使用 v3 模型
            handle_deepseek_command(&input[1..].trim(), "v3");
            continue;
        }
        
        // 直接执行普通命令
        if let Err(e) = execute_command_str(input) {
            eprintln!("执行错误: {}", e);
        }
    }
}

fn print_prompt() {
    let current_dir = env::current_dir().unwrap_or_else(|_| env::temp_dir());
    let username = whoami::username();
    let hostname = whoami::hostname();
    
    print!("{}@{}:{} $ ", username, hostname, current_dir.display());
}

fn execute_command(
    command: &str, 
    args: &[&str], 
    stdin_file: Option<&str>, 
    stdout_file: Option<&str>, 
    append: bool
) -> io::Result<()> {
    let mut cmd = Command::new(command);
    cmd.args(args);
    
    // 处理输入重定向
    if let Some(file) = stdin_file {
        let input_file = std::fs::File::open(file)?;
        cmd.stdin(input_file);
    } else {
        cmd.stdin(Stdio::inherit());
    }
    
    // 处理输出重定向
    if let Some(file) = stdout_file {
        let output_file = if append {
            OpenOptions::new().append(true).open(file)?
        } else {
            std::fs::File::create(file)?
        };
        cmd.stdout(output_file);
    } else {
        cmd.stdout(Stdio::inherit());
    }
    
    cmd.stderr(Stdio::inherit());
    
    let mut child = cmd.spawn()?;
    child.wait()?;
    Ok(())
}

fn execute_pipeline(input: &str) -> io::Result<()> {
    let commands: Vec<&str> = input.split('|').map(|s| s.trim()).collect();
    
    if commands.len() < 2 {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput, 
            "Pipeline requires at least two commands"
        ));
    }
    
    let mut previous_output: Option<std::process::ChildStdout> = None;
    
    for (i, cmd_str) in commands.iter().enumerate() {
        let cmd_parts: Vec<&str> = cmd_str.split_whitespace().collect();
        if cmd_parts.is_empty() {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput, 
                "Empty command in pipeline"
            ));
        }
        
        let cmd = cmd_parts[0];
        let args = &cmd_parts[1..];
        
        let mut current_cmd = Command::new(cmd);
        current_cmd.args(args);
        
        // 设置输入
        if i == 0 {
            // 第一个命令使用标准输入
            current_cmd.stdin(Stdio::inherit());
        } else if let Some(output) = previous_output.take() {
            // 中间命令使用前一个命令的输出作为输入
            current_cmd.stdin(output);
        } else {
            current_cmd.stdin(Stdio::inherit());
        }
        
        // 设置输出
        if i < commands.len() - 1 {
            // 中间命令的输出设置为管道
            current_cmd.stdout(Stdio::piped());
        } else {
            // 最后一个命令使用标准输出
            current_cmd.stdout(Stdio::inherit());
        }
        
        current_cmd.stderr(Stdio::inherit());
        
        let mut child = current_cmd.spawn()?;
        
        // 保存当前命令的输出以供下一个命令使用
        if i < commands.len() - 1 {
            previous_output = child.stdout.take();
        }
        
        // 等待当前命令完成
        child.wait()?;
    }
    
    Ok(())
}

fn execute_redirection(input: &str) -> io::Result<()> {
    let parts: Vec<&str> = input.split_whitespace().collect();
    let mut command_parts = Vec::new();
    let mut stdin_file = None;
    let mut stdout_file = None;
    let mut append = false;
    let mut i = 0;
    
    while i < parts.len() {
        match parts[i] {
            "<" => {
                if i + 1 < parts.len() {
                    stdin_file = Some(parts[i + 1]);
                    i += 1;
                } else {
                    return Err(io::Error::new(
                        io::ErrorKind::InvalidInput, 
                        "No file specified for input redirection"
                    ));
                }
            }
            ">" => {
                if i + 1 < parts.len() {
                    stdout_file = Some(parts[i + 1]);
                    append = false;
                    i += 1;
                } else {
                    return Err(io::Error::new(
                        io::ErrorKind::InvalidInput, 
                        "No file specified for output redirection"
                    ));
                }
            }
            ">>" => {
                if i + 1 < parts.len() {
                    stdout_file = Some(parts[i + 1]);
                    append = true;
                    i += 1;
                } else {
                    return Err(io::Error::new(
                        io::ErrorKind::InvalidInput, 
                        "No file specified for append redirection"
                    ));
                }
            }
            part => {
                command_parts.push(part);
            }
        }
        i += 1;
    }
    
    if command_parts.is_empty() {
        return Err(io::Error::new(
            io::ErrorKind::InvalidInput, 
            "No command specified"
        ));
    }
    
    let command = command_parts[0];
    let args = &command_parts[1..];
    
    execute_command(command, args, stdin_file, stdout_file, append)
}

// DeepSeek 命令建议功能
fn handle_deepseek_command(query: &str, model: &str) {
    // 调用 ask_deepseek 获取建议命令
    let model_name = match model {
        "v3" => "deepseek-chat",
        "r1" => "deepseek-reasoner",
        _ => "DeepSeek",
    };
    
    if let Some(suggested_command) = ask_deepseek(query, model_name) {
        println!("{} 建议执行: {}", model_name, suggested_command);
        
        // 询问用户是否要执行命令
        print!("是否执行此命令? (y/N): ");
        io::stdout().flush().unwrap();
        
        let mut response = String::new();
        if io::stdin().read_line(&mut response).is_err() {
            eprintln!("Error reading response");
            return;
        }
        
        // 检查用户确认
        let response = response.trim().to_lowercase();
        if response == "y" || response == "yes" {
            // 执行建议的命令
            if let Err(e) = execute_command_str(&suggested_command) {
                eprintln!("执行错误: {}", e);
            }
        } else {
            println!("取消执行");
        }
    }
    // 如果 ask_deepseek 返回 None，则不执行任何操作
}

// 执行命令字符串（支持管道和重定向）
fn execute_command_str(command_str: &str) -> io::Result<()> {
    // 特殊处理 cd 命令（必须在当前进程执行）
    if command_str.starts_with("cd ") {
        let dir = command_str.trim_start_matches("cd ").trim();
        if dir.is_empty() {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput, 
                "cd: missing argument"
            ));
        }
        
        return match env::set_current_dir(dir) {
            Ok(()) => Ok(()),
            Err(e) => Err(io::Error::new(
                io::ErrorKind::Other, 
                format!("cd: {}: {}", dir, e)
            )),
        };
    }
    
    // 检查是否包含管道或重定向
    if command_str.contains('|') {
        execute_pipeline(command_str)
    } else if command_str.contains('>') || command_str.contains('<') {
        execute_redirection(command_str)
    } else {
        // 普通命令
        let parts: Vec<&str> = command_str.split_whitespace().collect();
        if parts.is_empty() {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput, 
                "Empty command"
            ));
        }
        
        let command = parts[0];
        let args = &parts[1..];
        
        execute_command(command, args, None, None, false)
    }
}